import cv2
from model import recognizer
from model import recognizerTrainingModel
import ImageProcessingToolbox


def train(path):
    recognizerTrainingModel.TrainingModel.train(path)


def imageSegmentation(image_path):
    img = cv2.imread(image_path)

    # ============== 预处理 ================

    # 提取红色区域
    img = ImageProcessingToolbox.redPixelFilter(img)

    # 获取红色字体所在的区域
    img, contours = ImageProcessingToolbox.redFontScanner(img)
    print("sun = ", len(contours))

    # 获取矩形列表
    target_list, number_of_contours = ImageProcessingToolbox.getRectangleList(contours, img)
    print("number_of_contours = ", number_of_contours)

    # ================== 识别 =====================

    recognizeMachine = recognizer.Recognizer('model_plus.keras')

    # 遍历识别计算结果
    for i in range(len(target_list)):
        result, number = recognizeMachine.analyze(target_list[i].get_name())
        target_list[i].set_number(number)
        # 打印第几张图片，识别数字
        print(target_list[i].get_name(), " 识别数字为：", number)

    # ================== 后处理 ==================

    # 计算结果
    result_number = ImageProcessingToolbox.obtainTestScores(target_list)

    print("result_number = ", result_number)
    print("total score: ", 100 - result_number)

    # 显示图片
    cv2.namedWindow('img', cv2.WINDOW_NORMAL)
    cv2.imshow('img', img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == '__main__':
    imageSegmentation('./image/test2.png')

